package com.atbeijing.bigdata.spark.mytest.sql

import org.apache.spark.SparkConf
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object My_UDAF {
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udf")
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()
    import spark.implicits._

    val rdd: RDD[(Int, String, Int)] = spark.sparkContext.makeRDD(List(
      (1, "张三", 18),
      (1, "李四", 19),
      (1, "王五", 22),
      (1, "赵六", 30)
    ))
    //添加结构
    val df: DataFrame = rdd.toDF("id","name","age")
    df.createTempView("user")

    val myUDAF = new MyUDAF
    spark.udf.register("avgAge",myUDAF)
    spark.sql("select avgAge(age) from user").show()

  }

  //自定义UDAF(用户自定义聚合函数)弱类型
  class MyUDAF extends UserDefinedAggregateFunction {

    //输入参数类型
    override def inputSchema: StructType = {
      StructType(Array(StructField("age",IntegerType)))
    }

    //缓冲区中的数据类型
    override def bufferSchema: StructType = {
      StructType(Array(StructField("ageSum",IntegerType),StructField("ageCount",IntegerType)))
    }

    //函数返回值类型
    override def dataType: DataType = DoubleType

    // 稳定性：对于相同的输入是否一直返回相同的输出。
    override def deterministic: Boolean = true

    //缓冲区初始值
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      //ageSum
      buffer(0)=0
      //ageCount
      buffer(1)=0
    }

    //更新规则
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      if(!input.isNullAt(0)){
        buffer(0) = buffer.getInt(0)+input.getInt(0)
        buffer(1) = buffer.getInt(1)+1
      }
    }

    //合并分布式的结果
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1(0) = buffer1.getInt(0)+buffer2.getInt(0)
      buffer1(1) = buffer1.getInt(1)+buffer2.getInt(1)
    }

    //计算结果,这个buffer是最终结果
    override def evaluate(buffer: Row): Any = {
      buffer.getInt(0).toDouble/buffer.getInt(1)
    }

  }
}
